package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
)

func ensureKey(key string) []byte {
	result := key
	for i := 0; i < 32-len(key); i++ {
		result = "0" + result
	}
	return []byte(result)
}

func Encrypt(plain string, keyStr string) (string, error) {
	key := ensureKey(keyStr)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	plainBytes := []byte(plain)
	cipherText := make([]byte, len(plainBytes))
	stream := cipher.NewCTR(block, make([]byte, block.BlockSize()))
	stream.XORKeyStream(cipherText, plainBytes)
	return base64.StdEncoding.EncodeToString(cipherText), nil
}

func Decrypt(cipherText string, keyStr string) (string, error) {
	key := ensureKey(keyStr)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	cipherTextBytes, err := base64.StdEncoding.DecodeString(cipherText)
	if err != nil {
		return "", err
	}
	stream := cipher.NewCTR(block, make([]byte, block.BlockSize()))
	stream.XORKeyStream(cipherTextBytes, cipherTextBytes)
	return string(cipherTextBytes), nil
}
